{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tutorial 03: weak imposition of Dirichlet BCs by a Lagrange multiplier (linear problem)\n",
    "\n",
    "In this tutorial we solve the problem\n",
    "\n",
    "$$\\begin{cases}\n",
    "-\\Delta u = f, & \\text{in } \\Omega,\\\\\n",
    " u   = g, & \\text{on } \\Gamma = \\partial\\Omega,\n",
    "\\end{cases}$$\n",
    "\n",
    "where $\\Omega$ is the unit ball in 2D.\n",
    "\n",
    "We compare the following two cases:\n",
    "* **strong imposition of Dirichlet BCs**:\n",
    "the corresponding weak formulation is\n",
    "$$\n",
    "\\text{find } u \\in V_g \\text{ s.t. } \\int_\\Omega \\nabla u \\cdot \\nabla v = \\int_\\Omega f v, \\quad \\forall v \\in V_0\\\\\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V_g = \\{v \\in H^1(\\Omega): v|_\\Gamma = g\\},\\\\\n",
    "V_0 = \\{v \\in H^1(\\Omega): v|_\\Gamma = 0\\}.\\\\\n",
    "$$\n",
    "* **weak imposition of Dirichlet BCs**: this requires an introduction of a multiplier $\\lambda$ which is restricted to $\\Gamma$, and solves\n",
    "$$\n",
    "\\text{find } w, \\lambda \\in V \\times M \\text{ s.t. }\\\\\n",
    "\\begin{cases}\n",
    "\\int_\\Omega \\nabla w \\cdot \\nabla v + \\int_\\Gamma \\lambda v = \\int_\\Omega f v, & \\forall v \\in V,\\\\\n",
    "\\int_\\Gamma w \\mu = \\int_\\Gamma g \\mu, & \\forall \\mu \\in M\n",
    "\\end{cases}\n",
    "$$\n",
    "where\n",
    "$$\n",
    "V = H^1(\\Omega),\\\\\n",
    "M = L^{2}(\\Gamma).\\\\\n",
    "$$\n",
    "\n",
    "This example is a prototypical case of problems containing subdomain/boundary restricted variables (the Lagrange multiplier, in this case)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "from petsc4py import PETSc\n",
    "from ufl import grad, inner, Measure, TestFunction, TrialFunction\n",
    "from dolfinx import DirichletBC, Function, FunctionSpace\n",
    "from dolfinx.fem import LinearProblem, locate_dofs_topological\n",
    "from dolfinx.io import XDMFFile\n",
    "from dolfinx.plot import create_vtk_topology\n",
    "from multiphenicsx.fem import (assemble_matrix_block, assemble_scalar, assemble_vector_block,\n",
    "                               BlockVecSubVectorWrapper, create_vector_block, DofMapRestriction)\n",
    "import pyvista"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with XDMFFile(MPI.COMM_WORLD, \"data/circle.xdmf\", \"r\") as infile:\n",
    "    mesh = infile.read_mesh()\n",
    "    mesh.topology.create_connectivity_all()\n",
    "    subdomains = infile.read_meshtags(mesh, name=\"subdomains\")\n",
    "    boundaries = infile.read_meshtags(mesh, name=\"boundaries\")\n",
    "facets_Gamma = boundaries.indices[boundaries.values == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define associated measures\n",
    "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
    "ds = Measure(\"ds\")(subdomain_data=boundaries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dolfinx_to_pyvista_mesh(mesh):\n",
    "    num_cells = mesh.topology.index_map(mesh.topology.dim).size_local\n",
    "    cell_entities = np.arange(num_cells, dtype=np.int32)\n",
    "    pyvista_cells, cell_types = create_vtk_topology(mesh, mesh.topology.dim, cell_entities)\n",
    "    grid = pyvista.UnstructuredGrid(pyvista_cells, cell_types, mesh.geometry.x)\n",
    "    return grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_mesh_plot(mesh):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_mesh_plot(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Weak imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function space\n",
    "V = FunctionSpace(mesh, (\"Lagrange\", 2))\n",
    "M = V.clone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define restrictions\n",
    "dofs_V = np.arange(0, V.dofmap.index_map.size_local + V.dofmap.index_map.num_ghosts)\n",
    "dofs_M_Gamma = locate_dofs_topological(M, boundaries.dim, facets_Gamma)\n",
    "restriction_V = DofMapRestriction(V.dofmap, dofs_V)\n",
    "restriction_M_Gamma = DofMapRestriction(M.dofmap, dofs_M_Gamma)\n",
    "restriction = [restriction_V, restriction_M_Gamma]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define trial and test functions\n",
    "(u, l) = (TrialFunction(V), TrialFunction(M))\n",
    "(v, m) = (TestFunction(V), TestFunction(M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define problem block forms\n",
    "g = Function(V)\n",
    "g.interpolate(lambda x: np.sin(3 * x[0] + 1) * np.sin(3 * x[1] + 1))\n",
    "a = [[inner(grad(u), grad(v)) * dx, l * v * ds],\n",
    "     [u * m * ds, None]]\n",
    "f = [v * dx, g * m * ds]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Assemble the block linear system\n",
    "A = assemble_matrix_block(a, bcs=[], restriction=(restriction, restriction))\n",
    "A.assemble()\n",
    "F = assemble_vector_block(f, a, bcs=[], restriction=restriction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "ul = create_vector_block(f, restriction=restriction)\n",
    "ksp = PETSc.KSP()\n",
    "ksp.create(mesh.mpi_comm())\n",
    "ksp.setOperators(A)\n",
    "ksp.setType(\"preonly\")\n",
    "ksp.getPC().setType(\"lu\")\n",
    "ksp.getPC().setFactorSolverType(\"mumps\")\n",
    "ksp.setFromOptions()\n",
    "ksp.solve(F, ul)\n",
    "ul.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Split the block solution in components\n",
    "(u, l) = (Function(V), Function(M))\n",
    "with BlockVecSubVectorWrapper(ul, [V.dofmap, M.dofmap], restriction) as ul_wrapper:\n",
    "    for ul_wrapper_local, component in zip(ul_wrapper, (u, l)):\n",
    "        with component.vector.localForm() as component_local:\n",
    "            component_local[:] = ul_wrapper_local"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pyvista_scalar_field_plot(mesh, scalar_field, name):\n",
    "    grid = dolfinx_to_pyvista_mesh(mesh)\n",
    "    grid.point_arrays[name] = scalar_field.compute_point_values()\n",
    "    grid.set_active_scalars(name)\n",
    "    plotter = pyvista.PlotterITK()\n",
    "    plotter.add_mesh(grid)\n",
    "    plotter.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u, \"u\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, l, \"l\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Strong imposition of Dirichlet BCs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define Dirichlet BC object on Gamma\n",
    "dofs_V_Gamma = locate_dofs_topological(V, boundaries.dim, facets_Gamma)\n",
    "bc_ex = DirichletBC(g, dofs_V_Gamma)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solve\n",
    "u_ex = Function(V)\n",
    "problem_ex = LinearProblem(\n",
    "    a[0][0], f[0], bcs=[bc_ex], u=u_ex,\n",
    "    petsc_options={\"ksp_type\": \"preonly\", \"pc_type\": \"lu\", \"pc_factor_mat_solver_type\": \"mumps\"})\n",
    "problem_ex.solve()\n",
    "u_ex.vector.ghostUpdate(addv=PETSc.InsertMode.INSERT, mode=PETSc.ScatterMode.FORWARD)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyvista_scalar_field_plot(mesh, u_ex, \"u\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Comparison and error computation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "u_ex_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_ex), grad(u_ex)) * dx), op=MPI.SUM))\n",
    "err_norm = np.sqrt(mesh.mpi_comm().allreduce(assemble_scalar(inner(grad(u_ex - u), grad(u_ex - u)) * dx), op=MPI.SUM))\n",
    "print(\"Relative error is equal to\", err_norm / u_ex_norm)\n",
    "assert np.isclose(err_norm / u_ex_norm, 0., atol=1.e-10)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython"
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
